import pickle
import random
import sys

from Crypto.Cipher import AES
from Crypto.Random import get_random_bytes
import hashlib
import time

# 全局变量
g = 3
idA = b"User_A_8150"
idB = b"User_B_6623"
pw = b"ssx2016360."

sys.set_int_max_str_digits(20000)


# 建立哈希函数
def H(message):
    return int.from_bytes(hashlib.sha1(message).digest(), byteorder='big')


def H1(message):
    return int.from_bytes(hashlib.sha1(message).digest(), byteorder='big')


def H2(message):
    return int.from_bytes(hashlib.sha256(message).digest(), byteorder='big')


def H3(message):
    return int.from_bytes(hashlib.sha1(message).digest(), byteorder='big')


# 建立加密函数
def encrypt(key, plaintext):
    cipher = AES.new(key, AES.MODE_GCM)
    nonce = cipher.nonce  # 保证加密和解密是同一个会话，唯一的
    ciphertext, tag = cipher.encrypt_and_digest(plaintext)
    return ciphertext, nonce, tag


# 建立解密函数
def decrypt(key, ciphertext, nonce, tag):
    cipher = AES.new(key, AES.MODE_GCM, nonce=nonce)
    plaintext = cipher.decrypt(ciphertext)
    try:
        cipher.verify(tag)
        return plaintext
        # print(plaintext)
    except ValueError:
        print("key is not valid")


# 建立时间戳验证函数
def verify_timestamp(timestamp):
    current_time = time.time()  # 获取当前时间戳
    if current_time - timestamp > 10:  # 判断时间戳之间的差距，单位为秒
        return False
    else:
        return True


def generate_password_file(pw, idd):
    s = int.from_bytes(get_random_bytes(10), byteorder="big")  # 将随机生成的随机数从字节转换成int
    hw = H1(pw) ^ s
    D = H2(s.to_bytes(s.bit_length(), byteorder='big')) ^ int.from_bytes(idd, byteorder='big')
    return hw, D


def session():
    # UA部分，生成XA、YA、TA发送给UB
    rA = int.from_bytes(get_random_bytes(1), byteorder="big")
    hwA, DA = generate_password_file(pw, idA)
    XA = H2(hwA.to_bytes(hwA.bit_length(), byteorder='big')) ^ rA
    YA = g ** rA
    TA = time.time()
    # UB部分，收到XA、YA、TA之后
    rB = int.from_bytes(get_random_bytes(1), byteorder="big")
    hwB, DB = generate_password_file(pw, idB)
    XB = H2(hwB.to_bytes(hwB.bit_length(), byteorder='big')) ^ rB
    YB = g ** rB
    if not verify_timestamp(TA):
        print("TA is not fresh")
    else:
        TB = time.time()
        tkB = XA ^ DB ^ int.from_bytes(idB, byteorder='big') ^ H2(rB.to_bytes(rB.bit_length(), byteorder='big')) ^ H2(
            H3((YA ** rB).to_bytes((YA ** rB).bit_length(), byteorder='big')).to_bytes(
                H3((YA ** rB).to_bytes((YA ** rB).bit_length(), byteorder='big')).bit_length(), byteorder='big'))
        CB, nonce_b, tag_b = encrypt(tkB, pickle.dumps(TB, hwB, rB))  # nonce和tag是加密的唯一对，解密需要验证，dumps是将元组转换成字节
        # UB运行完毕，将XB、YB、TB、CB发送给UA
        if not verify_timestamp(TB):
            print("TB is not fresh")
        else:
            tkA = XB ^ DA ^ int.from_bytes(idA, byteorder='big') ^ H2(rA) ^ H2(H3(YB ** rA))
            plaintext_b = pickle.loads(decrypt(tkA, CB, nonce_b, tag_b))  # 将字节转换成元组
            dec_TB = plaintext_b[0]
            dec_hwB = plaintext_b[1]
            dec_rB = plaintext_b[2]
            if dec_TB != TB:
                print("TB和dec_TB时间戳不一致！")
            else:
                if YB != g ** dec_rB or XB != H2(dec_hwB) ^ H2(dec_rB):
                    print("YB或XB不一致！")
                else:
                    CA, nonce_a, tag_a = encrypt(tkA, (TA, hwA, rA))
                    sidA = XA + YA + XB + YB
                    K = H(tkA + sidA)
                    print("用户A的密钥Ka=" + K)
                    # UA处理结束，将CA加密发送给UB
                    plaintext_a = decrypt(tkB, CA, nonce_a, tag_a)
                    dec_TA = plaintext_a[0]
                    dec_hwA = plaintext_a[1]
                    dec_rA = plaintext_a[2]
                    if dec_TA != TA:
                        print("TA和dec_TA时间戳不一致！")
                    else:
                        if YA != g ** dec_rA or XA != H2(dec_hwA) ^ H2(dec_rA):
                            print("YA或XA不一致！")
                        else:
                            sidB = XA + YA + XB + YB
                            K = H(tkB + sidB)
                            print("用户B的密钥Ka=" + K)


if __name__ == '__main__':
    # print(type(H(pw)), type(H1(pw)), type(H2(pw)), type(H3(pw)), type(get_random_bytes(32)))
    # a = int(bin(234567890))
    # a = 23456789
    # b = 98765432
    # b = int(hex(34567890))
    # print(a, type(a), b, type(b), a ^ b)
    # e = get_random_bytes(32)
    # int_e = int.from_bytes(e, byteorder='big')
    # print(e, type(e), int_e, type(int_e))
    session()
    # print(H2(9876543234567890987654356753656536646354568523466325645634566234562345655345436566756654346566364565364))
    # s = get_random_bytes(10)
    # int_s = int.from_bytes(s, byteorder="big")  # 将随机生成的随机数从字节转换成int
    # byte_s = int_s.to_bytes(int_s.bit_length(), byteorder="big")
    # sss = int.from_bytes(byte_s, byteorder="big")
    # print(s)
    # print(int_s)
    # print(byte_s)
    # print(sss)

    # print(3**15816)
